##################################################
# (1) init ==> B <= 0
# (2) unsafe ==> B > 0
# (3) domain /\ x /\ x' ==> B(x') - (1-alpha) * B(x) < 0
##################################################
# max(B(x))  s.t. x in init
# if max(B)<=0: condition (1) is OK
import math
import time
import cvxpy as cp
import dccp
import gurobipy as gp
from gurobipy import GRB

import verify_dccp.getMathExpress2 as express
from prob import *
import numpy as np
import verify_piecewise_utils
from NN import cdinn2
import torch


def verif(model, piece_len):
    cex_init = []
    res_init = []
    cex_unsafe = []
    res_unsafe = []
    cex_domain = []
    res_domain = []

    x = cp.Variable(DIM)  # x_0, x_1
    v = cp.Variable()

    f1, f2 = express.getExpr(x, model)

    # print("curvature of f1:", f1.curvature)
    # print("curvature of f2:", f2.curvature)

    # print("=====================X_INIT=========================")
    objective = cp.Maximize(f1 - v)
    constraints = [v == f2]

    constraints += [cp.square(x[0]) + cp.square(x[1] - 1) <= 0.25]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)  # [ 1.34102711 -0.47405446]
    # print("v.value:", v.value)  # 4.158081062071981
    # print("status:", problem.status) # Converged
    # print("cost value0 =", result[0])  # -1.161325675958789
    if result[0] > 0.0:
        cex_init.append(x.value)
        res_init.append(result[0])
    else:
        print("init is ok")

    # print("=====================X_UNSAFE=========================")
    objective = cp.Minimize(v - f2)
    constraints = [v == f1]

    constraints += [cp.square(x[0] - 1.5) + cp.square(x[1] - 1.5) <= 0.25]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    # x.value = np.array([1, 2])
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)
    # print("v.value:", v.value)
    # print("status:", problem.status)
    # print("cost value0 =", result[0])
    if result[0] < 0.0:
        cex_unsafe.append(x.value)
        res_unsafe.append(result[0])
    else:
        print("unsafe is ok")

    # print("=====================X_DOMAIN=========================")
    # fw = open('log/verify_eg4_0-01.txt', 'w+', encoding='utf-8')

    data_dir = "./data/"
    pieces = int(4 / piece_len)
    # datas = np.loadtxt(data_dir + "10points_0.1.txt", dtype=np.float64).reshape([-1, 8])
    # datas = np.loadtxt(data_dir + "-2_2_0.5_20_new.txt", dtype=np.float64).reshape([-1, 8])
    datas = np.loadtxt(data_dir + f"21_-2_2_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 4])
    new_cex_idx_list = []

    var_b = cp.Variable()

    x0_min = domain_min[0]  # -2
    x0_max = domain_max[0]  # 2
    x1_min = domain_min[1]  # -2
    x1_max = domain_max[1]  # 2

    cex_idx_list = []

    if len(cex_idx_list) == 0:
        for x0_idx in range(pieces):
            for x1_idx in range(pieces):
                cex_idx = [x0_idx, x1_idx]
                cex_idx_list.append(cex_idx)

    for cex_idx in cex_idx_list:
        x0_idx, x1_idx = cex_idx

        # w0, w1, b_l, b = datas[x0_idx * pieces + x1_idx - 1]
        w0, w1, b_l, b = datas[x0_idx * pieces + x1_idx]

        x0_l = x0_idx * piece_len + x0_min
        x0_r = (x0_idx + 1) * piece_len + x0_min
        x1_l = x1_idx * piece_len + x1_min
        x1_r = (x1_idx + 1) * piece_len + x1_min

        x0_next = time_step * (- x[0] + 2 * w0 * x[0] + w1 * x[1] + b_l + var_b) + x[0]
        x1_next = time_step * (- x[1]) + x[1]

        f1_next, f2_next = express.getExpr([x0_next, x1_next], model)

        # print("curvature of f1_next:", f1_next.curvature)
        # print("curvature of f2_next:", f2_next.curvature)

        u = cp.Variable()
        objective = cp.Maximize((f1_next - v) - (1 - alpha) * (u - f2))
        constraints = [v == f2_next, u == f1]
        constraints += [x0_l <= x[0], x[0] <= x0_r]
        constraints += [x1_l <= x[1], x[1] <= x1_r]
        constraints += [x0_min <= x0_next, x0_next <= x0_max]
        constraints += [x1_min <= x1_next, x1_next <= x1_max]

        constraints += [0 <= var_b, var_b <= b]

        problem = cp.Problem(objective, constraints)
        # print("problem is DCCP:", dccp.is_dccp(problem))
        # x.value = np.array([1, 2])
        result = problem.solve(method='dccp')
        # print("x.value:", x.value)
        # print("var_b.value:", var_b.value)
        # print("problem.status:", problem.status)
        if problem.status == "Converged":
            # print("cost value =", result[0])
            # fw.write("================[" + str(x_l) + "," + str(x_r) + "]===================")
            # fw.write('\n')
            # fw.write(str(result[0]))
            # fw.write('\n')
            # fw.write('\n')

            if result[0] > 0.0:
                cex_domain.append(x.value)
                res_domain.append(result[0])
                # print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}] has cex point")
                # new_cex_idx_list.append(cex_idx)
        #     else:
        #         print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}] is ok")
        # else:
        #     print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}] is not converged")

    # fw.close()
    # print(new_cex_idx_list)
    # print("==============cex and res=============")
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../model/cdinn_c13_train.pt'), strict=True)
    # time_start_verify = time.time()
    verif(model)
    # time_end_verify = time.time()
    # print("verifing totally costs:", time_end_verify - time_start_verify)



